from flask import Flask, jsonify, request
from flask_cors import CORS
from futu import *
import pandas as pd
from datetime import datetime

app = Flask(__name__)
CORS(app)

quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)

# 在程序启动时订阅需要的 K 线类型
def ensure_subscribe(stock, ktypes):
    for kt in ktypes:
        ret, err = quote_ctx.subscribe([stock], [kt], subscribe_push=False)
        if ret != RET_OK:
            print("订阅失败:", kt, err)

ensure_subscribe('HK.02616', [SubType.K_1M, SubType.K_5M, SubType.K_15M, SubType.K_30M,
                              SubType.K_60M, SubType.K_DAY, SubType.K_WEEK, SubType.K_MON])

@app.route('/kline')
def get_kline():
    stock = request.args.get('code', 'HK.02616')
    ktype_str = request.args.get('type', '1m').lower()
    num = int(request.args.get('num', '500'))
    if num > 1000:
        num = 1000

    ktype_map = {
        '1m': KLType.K_1M,
        '5m': KLType.K_5M,
        '15m': KLType.K_15M,
        '30m': KLType.K_30M,
        '1h': KLType.K_60M,
        'day': KLType.K_DAY,
        'week': KLType.K_WEEK,
        'month': KLType.K_MON
    }
    ktype = ktype_map.get(ktype_str, KLType.K_1M)
    autype = AuType.QFQ

    ret, data = quote_ctx.get_cur_kline(stock, num=num, ktype=ktype, autype=autype)
    if ret != RET_OK:
        return jsonify({"error": data})

    # 确保数据是 DataFrame
    if not isinstance(data, pd.DataFrame) or data.empty:
        return jsonify({"error": "无数据"})

    # 过滤当天的数据（仅对分钟级别有效）
    if ktype in [KLType.K_1M, KLType.K_5M, KLType.K_15M, KLType.K_30M, KLType.K_60M]:
        today_str = datetime.now().strftime('%Y-%m-%d')
        data = data[data['time_key'].astype(str).str.startswith(today_str)]

    times = data['time_key'].astype(str).tolist()
    opens = data['open'].tolist()
    closes = data['close'].tolist()
    highs = data['high'].tolist()
    lows = data['low'].tolist()
    volumes = data['volume'].fillna(0).tolist()
    turnovers = data['turnover'].fillna(0).tolist()
    turnover_rates = data.get('turnover_rate', None)
    if turnover_rates is not None:
        turnover_rates = turnover_rates.fillna(0).tolist()
    else:
        turnover_rates = [0] * len(times)

    current_price = closes[-1] if closes else None
    highest_price = max(highs) if highs else None
    lowest_price = min(lows) if lows else None

    return jsonify({
        "code": stock,
        "times": times,
        "kline": list(map(list, zip(opens, closes, lows, highs))),
        "volumes": volumes,
        "turnovers": turnovers,
        "turnover_rates": turnover_rates,
        "current_price": current_price,
        "highest_price": highest_price,
        "lowest_price": lowest_price
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
